# generate by ollama qwen:8b

import pygame
import random
import time

# 初始化 Pygame
pygame.init()

# 屏幕尺寸
WIDTH, HEIGHT = 600, 400
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Snake Game")

# 颜色定义
WHITE = (255, 255, 250)
GREEN = (0, 255, 0)
RED = (255, 0, 0)
BLACK = (0, 0, 0)

# 字体设置
FONT = pygame.font.SysFont("Arial", 24)

# 游戏参数
BLOCK_SIZE = 20
SPEED = 15


class Snake:
    def __init__(self):
        self.body = [(WIDTH // 2, HEIGHT // 2)]
        self.direction = "RIGHT"

    def move(self):
        x, y = self.body[0]
        if self.direction == "UP":
            new_head = (x, y - BLOCK_SIZE)
        elif self.direction == "DOWN":
            new_head = (x, y + BLOCK_SIZE)
        elif self.direction == "LEFT":
            new_head = (x - BLOCK_SIZE, y)
        else:  # RIGHT
            new_head = (x + BLOCK_SIZE, y)

        self.body.insert(0, new_head)

    def check_collision(self):
        x, y = self.body[0]
        if (x < 0 or x >= WIDTH or
                y < 0 or y >= HEIGHT or
                (x, y) in self.body[1:]):
            return True
        return False

    def grow(self):
        self.body.append(self.body[-1])


class Food:
    def __init__(self):
        self.position = self.generate_position()

    def generate_position(self):
        while True:
            x = random.randint(0, (WIDTH - BLOCK_SIZE) // BLOCK_SIZE) * BLOCK_SIZE
            y = random.randint(0, (HEIGHT - BLOCK_SIZE) // BLOCK_SIZE) * BLOCK_SIZE
            if (x, y) not in snake.body:
                return (x, y)

    def draw(self):
        pygame.draw.rect(screen, RED, (*self.position, BLOCK_SIZE, BLOCK_SIZE))


def draw_score(score):
    score_text = FONT.render(f"Score: {score}", True, WHITE)
    screen.blit(score_text, (10, 10))


def game_over():
    screen.fill(BLACK)
    game_over_text = FONT.render("Game Over", True, RED)
    screen.blit(game_over_text, (WIDTH // 2 - 80, HEIGHT // 2 - 20))
    pygame.display.flip()
    time.sleep(2)
    pygame.quit()
    quit()


# 初始化游戏对象
snake = Snake()
food = Food()
score = 0
clock = pygame.time.Clock()

# 主游戏循环
running = True
while running:
    clock.tick(SPEED)

    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False

        # 控制方向
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_UP and snake.direction != "DOWN":
                snake.direction = "UP"
            elif event.key == pygame.K_DOWN and snake.direction != "UP":
                snake.direction = "DOWN"
            elif event.key == pygame.K_LEFT and snake.direction != "RIGHT":
                snake.direction = "LEFT"
            elif event.key == pygame.K_RIGHT and snake.direction != "LEFT":
                snake.direction = "RIGHT"

    # 更新游戏状态
    snake.move()

    # 检查碰撞
    if snake.check_collision():
        game_over()

    # 吃食物
    if snake.body[0] == food.position:
        score += 1
        food = Food()
    else:
        snake.body.pop()

    # 绘制画面
    screen.fill(BLACK)

    # 绘制蛇
    for segment in snake.body:
        pygame.draw.rect(screen, GREEN, (*segment, BLOCK_SIZE, BLOCK_SIZE))

    # 绘制食物
    food.draw()

    # 绘制得分
    draw_score(score)

    pygame.display.flip()

pygame.quit()
